#!/usr/bin/env python
#coding:utf-8

# 识别物体颜色
from ctypes import POINTER
import rospy
import cv2
from sensor_msgs.msg import CompressedImage
from geometry_msgs.msg import Pose
from cv_bridge import CvBridge
import numpy as np

# ------------------------------参数集合------------------------------
# 颜色HSV范围
colors = {'red': {'Lower': np.array([156, 43, 46]), 'Upper': np.array([180, 255, 255])},
          'yellow': {'Lower': np.array([26, 43, 46]), 'Upper': np.array([34, 255, 255])},
          'green': {'Lower': np.array([35, 43, 46]), 'Upper': np.array([85, 255, 255])},
          'blue': {'Lower': np.array([84, 43, 46]), 'Upper': np.array([107, 255, 255])}}

# 控制点在世界坐标系的坐标
objPoints = np.array([[0, 0, 0],
                      [0, 20, 0],
                      [20, 20, 0],
                      [20, 0, 0]], dtype=np.float64)

# 相机内参矩阵与外参矩阵
cameraMatrix = np.array([[664.763784, 0, 333.286264],
                         [0, 663.302322, 215.565802],
                         [0, 0, 1]], dtype=np.float64)
distCoeffs = np.array([-0.354213, 0.167628, -0.000365, -0.001693, 0], dtype=np.float64)

# 颜色转数字
color_num = {'red':1, 'yellow':2, 'green':3, 'blue':4}

pnp_msg = Pose()


# ------------------------------识别颜色函数------------------------------
def color(frame, pnp_vel_pub):
	hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)     # 转化成HSV图像

	for (k, (ball_color, rgb)) in enumerate(colors.items()):  # 提取一种颜色的HSV范围
		ranges = cv2.inRange(hsv, colors[ball_color]['Lower'], colors[ball_color]['Upper'])
		gs_frame = cv2.GaussianBlur(ranges, (5, 5), 0)      # 高斯模糊
		erode_hsv = cv2.erode(gs_frame, None, iterations=2)    # 先腐蚀再膨胀
		dilate = cv2.dilate(erode_hsv.copy(), None, iterations=2)

		cnts = cv2.findContours(dilate.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
		if cnts:
			frame = point(cnts, frame, ball_color, pnp_vel_pub)
	return frame


# ------------------------------拟合轮廓寻找角点函数------------------------------
def point(cnts, frame, ball_color, pnp_vel_pub):
	approx = []
	# c = max(cnts, key=cv2.contourArea)
	num = 0
	for j in range(len(cnts)):
		((x, y), radius) = cv2.minEnclosingCircle(cnts[j])  # 找到最小的圆包含这个轮廓，返回坐标和半径
		if radius > 20:  # 判断轮廓的半径，太小的话就认为是噪声，就忽略掉
			approx = cv2.approxPolyDP(cnts[j], 40, True)  # 拟合轮廓得到角点
			if len(approx) == 4:
				num += 1
				for i in range(len(approx)):
					frame = cv2.line(frame, (approx[i][0][0], approx[i][0][1]),
								 (approx[(i + 1) % 4][0][0], approx[(i + 1) % 4][0][1]), (255, 0, 0), 3)
					cv2.circle(frame, (approx[i][0][0], approx[i][0][1]), 4, (0, 0, 255), -1)
						
				# w 角点颜色	x 角点数目
				pnp_msg.orientation.w = color_num[ball_color]
				# pnp_msg.orientation.x = num

				# 发布消息
				# pnp_vel_pub.publish(pnp_msg)
				rospy.loginfo("%s", ball_color)
				rospy.loginfo("%d", num)

				# pnp测距
				pnp(approx, pnp_vel_pub)
	return frame


# ------------------------------pnp算法测距函数------------------------------
def pnp(approx, pnp_vel_pub):
	# 检测到的角点的坐标
	imgPoints = np.array([[approx[0][0][0], approx[0][0][1]],
		          [approx[1][0][0], approx[1][0][1]],
		          [approx[2][0][0], approx[2][0][1]],
		          [approx[3][0][0], approx[3][0][1]]], dtype=np.float64)
	retval, rvec, tvec = cv2.solvePnP(objPoints, imgPoints, cameraMatrix, distCoeffs)
	R = cv2.Rodrigues(rvec)
	R = np.transpose(R[0])
	P = np.transpose(tvec)

	# 初始化geometry_msgs::Pose类型的消息
	pnp_msg.position.x = P[0][0]
	pnp_msg.position.y = P[0][1]
	pnp_msg.position.z = P[0][2]
	rospy.loginfo("%s", P[0][0])
	rospy.loginfo("%s", P[0][1])
	rospy.loginfo("%s", P[0][2])

	# 发布消息
	pnp_vel_pub.publish(pnp_msg)


def main():

	# ROS节点初始化
	rospy.init_node('polygon_publisher', anonymous=True)
	
	# 创建两个Publisher，发布名为/vision/camera/image_raw的topic，消息类型为sensor_msgs::Image，队列长度10
	img_raw_pub = rospy.Publisher('/vision/camera/image_raw',CompressedImage, queue_size=1)
	img_hd_pub = rospy.Publisher('/vision/camera/image_handle',CompressedImage, queue_size=1)
	
	# 创建一个Publisher，发布名为/vision/cmd/gripper的topic，消息类型为geometry_msgs::Pose，队列长度10
	pnp_vel_pub = rospy.Publisher('/vision/cmd/gripper', Pose, queue_size=1)

	bridge = CvBridge()

	# 创建压缩图像
	msg = CompressedImage()

	cap = cv2.VideoCapture(2)
	rate = rospy.Rate(10)
	while not rospy.is_shutdown():

		# 获取相机图像
		ret, frame = cap.read()

		# 发布原始图像
		msg.header.stamp = rospy.Time.now()
		msg.format = "jpeg"
		msg.data = np.array(cv2.imencode('.jpg', frame)[1]).tostring()
		img_raw_pub.publish(msg)


		frame2 = color(frame, pnp_vel_pub)
		
		# 结束信息
		# pnp_msg = Pose()
		# pnp_msg.orientation.y = 0
		# pnp_vel_pub.publish(pnp_msg)
		# rospy.loginfo("%s", pnp_msg.orientation.y)
		
		# 发布处理后的图像
		msg.header.stamp = rospy.Time.now()
		msg.format = "jpeg"
		msg.data = np.array(cv2.imencode('.jpg', frame2)[1]).tostring()
		img_hd_pub.publish(msg)

		# cv2.imshow('photo',frame2)
		# cv2.waitKey(10)

		rate.sleep()

if __name__ == "__main__":
	try:
		main()
	except rospy.ROSInterruptException:
		pass


